package com.fenquen.sourceguard.repack;

import com.fenquen.sourceguard.dencrypt.Dencrypt;
import com.fenquen.sourceguard.repack.util.StringUtils;

import java.io.*;
import java.util.*;
import java.util.concurrent.*;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;
import java.util.jar.JarOutputStream;
import java.util.zip.CRC32;
import java.util.zip.ZipEntry;

public class Repackager {

    private static ExecutorService POOL = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors());

    private static Vector<Future<Void>> FUTURE_LIST = new Vector<>(1024);

    private static final String SUFFIX = ".jar";

    public static boolean bindMachine = false;

    /**
     * @param originJarPath            待加密的jar文件的path
     * @param packageNameClassNameList 需要加密的package/class名字
     */
    public static void encryptJar(String originJarPath,
                                  List<String> packageNameClassNameList) throws Exception {
        String encryptedJarPath;

        // 要是以".jar"结尾例如a.jar,那么生成的jar命名为a_encrypted.jar
        if (originJarPath.endsWith(SUFFIX)) {
            encryptedJarPath = originJarPath.substring(0, originJarPath.length() - SUFFIX.length()) + "_encrypted.jar";
        } else { // 要是不以".jar"结尾直接在末尾添加"_encrypted.jar"
            encryptedJarPath = originJarPath + "_encrypted.jar";
        }

        encryptJar(originJarPath, encryptedJarPath, packageNameClassNameList);
    }

    /**
     * @param originJarPath       待加密的jar文件的path
     * @param encryptedJarPath    生成的加密jar文件的path
     * @param basePackageNameList 需要加密的类对应的包
     */
    public static void encryptJar(String originJarPath,
                                  String encryptedJarPath,
                                  List<String> basePackageNameList) throws Exception {
        // 用来在解压jar文件时候记录 隶属的嵌套的子jar文件在原始根jar文件中的path->解压后的jarEntry的在硬盘absPath->jarEntry
        Map<Pair, Map<String, JarEntry>> subJar_absPath_entryName = new ConcurrentHashMap<>(1024);

        // jar文件解压之后保存内容的目录,和jar文件在同级目录
        String unpackedJarDirPath = "./unpacked";
        unpackJar2Dir(originJarPath, unpackedJarDirPath, subJar_absPath_entryName, Pair.EMPTY);

        waitDone(FUTURE_LIST);

        // 要把"."替换为"/"
        for (int a = 0; basePackageNameList.size() > a; a++) {
            basePackageNameList.set(a, basePackageNameList.get(a).replaceAll("\\.", "/"));
        }

        try (JarOutputStream encryptedJarOutputStream = new JarOutputStream(new FileOutputStream(encryptedJarPath))) {
            repack(subJar_absPath_entryName, encryptedJarOutputStream, basePackageNameList);
            waitDone(FUTURE_LIST);
        }

        rm("./unpacked");

        rm("./repacked_sub_jar");

        List<Runnable> runnableList = POOL.shutdownNow();
        assert runnableList.size() == 0;
    }

    /**
     * 把待加密的jar文件解压到指定目录
     *
     * @param originJarPath            待加密的jar文件/内嵌jar文件的path
     * @param unpackedJarDirPath       解压目录的path
     * @param subJar_absPath_entryName subJar->absPath->entryName
     * @param parentPair               当前对应的嵌套的jar,要是当前在根jar维度那么为空之类
     */
    private static void unpackJar2Dir(final String originJarPath,
                                      final String unpackedJarDirPath,
                                      final Map<Pair, Map<String, JarEntry>> subJar_absPath_entryName,
                                      final Pair parentPair) throws Exception {
        // 说明之前是在嵌套的jar包,但是又到了个嵌套的jar,相当是嵌套又嵌套jar,直接不允许
        if (parentPair.subJarGlobalEntryName.contains("@")) {
            throw new UnsupportedOperationException("嵌套的jar中又嵌套了jar");
        }

        File unpackedJarDir = new File(unpackedJarDirPath);

        if (!unpackedJarDir.exists() && !unpackedJarDir.mkdirs()) {
            throw new RuntimeException("无法创建解压目录: " + unpackedJarDirPath);
        }

        JarFile originJarFile = new JarFile(originJarPath);

        Enumeration<JarEntry> entries = originJarFile.entries();
        while (entries.hasMoreElements()) {
            JarEntry jarEntry = entries.nextElement();

            //System.out.println(jarEntry.getName());

            File destFile = new File(unpackedJarDir, jarEntry.getName());
            String destFileAbsPath = destFile.getAbsolutePath();

            // 目录,不能跳过在重新生成jar的时候需要显式调用
            if (jarEntry.isDirectory()) {
                Map<String, JarEntry> absPath_entryName = subJar_absPath_entryName.computeIfAbsent(parentPair, key -> new HashMap<>());
                absPath_entryName.put(destFileAbsPath, jarEntry);
                continue;
            }

            // System.out.println(jarEntry.getMethod() + " " + jarEntry.getSize() + " " + jarEntry.getCompressedSize() + " " + jarEntry.getCrc());

            // 把当前的jarEntry(不管是普通的还是内嵌jar)拷贝到unpackedJarDirPath
            copyJarEntry2Dest(originJarFile, jarEntry, destFile);

            String prefix = parentPair.subJarGlobalEntryName;
            // 说明是已在递归的,当前对应的jar是一个嵌套的不是根
            if (StringUtils.isNotBlank(parentPair.subJarGlobalEntryName)) {
                prefix += "@";
            }

            // 当前的jarEntry是内嵌的jar文件,单线程时候是递归调用多线程时候使用独立线程
            if (destFile.getName().endsWith(SUFFIX)) {
                String prefix_ = prefix;

                Future<Void> unpackSubJarFuture = POOL.submit(() -> {
                    unpackJar2Dir(destFileAbsPath,
                            unpackedJarDirPath + "/" + jarEntry.getName() + "_unpacked",
                            subJar_absPath_entryName, new Pair(prefix_ + jarEntry.getName(), jarEntry));

                    destFile.delete();

                    return null;
                });

                FUTURE_LIST.add(unpackSubJarFuture);
            } else {
                Map<String, JarEntry> absPath_entryName = subJar_absPath_entryName.computeIfAbsent(parentPair, key -> new ConcurrentHashMap<>());
                absPath_entryName.put(destFileAbsPath, jarEntry);
            }
        }

    }

    /**
     * 拷贝jarEntry到指定地点
     *
     * @param jarFile  该jarEntry隶属的jar
     * @param jarEntry 该jarEntry
     * @param destFile
     */
    private static void copyJarEntry2Dest(JarFile jarFile,
                                          JarEntry jarEntry,
                                          File destFile) throws Exception {
        File parentDir = destFile.getParentFile();
        if (!parentDir.exists() && !parentDir.mkdirs()) {
            throw new IOException("");
        }

        try (InputStream inputStream = jarFile.getInputStream(jarEntry);
             OutputStream outputStream = new FileOutputStream(destFile)) {
            byte[] buffer = new byte[1024];
            int bytesRead;
            while ((bytesRead = inputStream.read(buffer)) != -1) {
                outputStream.write(buffer, 0, bytesRead);
            }
            outputStream.flush();
        }
    }

    public static void repack(final Map<Pair, Map<String, JarEntry>> subJar_absPath_entryName,
                              JarOutputStream encryptedJar,
                              List<String> basePackageNameList) throws Exception {

        File repackedSubJarParentDir = new File("./repacked_sub_jar");
        repackedSubJarParentDir.mkdirs();

        for (Map.Entry<Pair, Map<String, JarEntry>> entry : subJar_absPath_entryName.entrySet()) {
            // lib/a.jar
            String subJarGlobalEntryName = entry.getKey().subJarGlobalEntryName;
            Map<String, JarEntry> absPath_entry = entry.getValue();

            // 说明是在嵌套jar维度需要先把该嵌套的jar生成好,多线程模式下使用独立线程
            if (StringUtils.isNotBlank(subJarGlobalEntryName)) {
                // 重打包好的内嵌jar的path
                File repackedSubJar = new File(repackedSubJarParentDir, UUID.randomUUID() + "_repacked.jar");

                Future<Void> reassembleAndPutSubJarFuture = POOL.submit(() -> {
                    //System.out.println(subJarGlobalEntryName);
                    try (JarOutputStream encryptedJarSub = new JarOutputStream(new FileOutputStream(repackedSubJar))) {
                        //encryptedJar.setMethod(ZipOutputStream.STORED);
                        Map<Pair, Map<String, JarEntry>> subJar_absPath_entryName4Sub = new HashMap<>();
                        subJar_absPath_entryName4Sub.put(Pair.EMPTY, absPath_entry);

                        // encryptedJar.setMethod(entry.getKey().jarEntry.getMethod());

                        repack(subJar_absPath_entryName4Sub, encryptedJarSub, basePackageNameList);
                    }

                    // 把生成的嵌套jar文件注入到根jar文件
                    try (ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
                         FileInputStream fileInputStream = new FileInputStream(repackedSubJar)) {

                        byte[] buf = new byte[1024];

                        int len;

                        while ((len = fileInputStream.read(buf, 0, buf.length)) != -1) {
                            byteArrayOutputStream.write(buf, 0, len);
                        }

                        byte[] bytes = byteArrayOutputStream.toByteArray();

                        // 内嵌jar的压缩方式要求是store并不压缩
                        // JarEntry jarEntry = new JarEntry(entry.getKey().jarEntry);
                        JarEntry jarEntry = new JarEntry(entry.getKey().subJarGlobalEntryName);

                        jarEntry.setMethod(ZipEntry.STORED);
                        jarEntry.setSize(bytes.length);
                        jarEntry.setCompressedSize(bytes.length);
                        CRC32 crc32 = new CRC32();
                        crc32.update(bytes);
                        jarEntry.setCrc(crc32.getValue());

                        synchronized (encryptedJar) {
                            encryptedJar.putNextEntry(jarEntry);
                            encryptedJar.write(bytes);
                            encryptedJar.closeEntry();
                        }

                        // countDownLatch.countDown();
                        byteArrayOutputStream.reset();
                    }

                    return null;
                });

                FUTURE_LIST.add(reassembleAndPutSubJarFuture);

                continue;
            }

            // 不是在嵌套的jar维度
            try (ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream()) {
                byte[] buf = new byte[1024];

                for (Map.Entry<String, JarEntry> entry_ : absPath_entry.entrySet()) {
                    String absPath = entry_.getKey();
                    String entryName = entry_.getValue().getName();

                    // 目录也要显式创建
                    if (entry_.getValue().isDirectory()) {
                        JarEntry jarEntry = new JarEntry(entry_.getValue().getName());

                        jarEntry.setMethod(ZipEntry.STORED);
                        jarEntry.setSize(0);
                        jarEntry.setCompressedSize(0);
                        jarEntry.setCrc(0);

                        synchronized (encryptedJar) {
                            encryptedJar.putNextEntry(jarEntry);
                            encryptedJar.closeEntry();
                        }
                        continue;
                    }

                    // 普通的
                    try (FileInputStream fileInputStream = new FileInputStream(absPath)) {
                        int len;

                        while ((len = fileInputStream.read(buf, 0, buf.length)) != -1) {
                            byteArrayOutputStream.write(buf, 0, len);
                        }

                        byte[] bytes = byteArrayOutputStream.toByteArray();

                        for (String basePackageName : basePackageNameList) {
                            if (entryName.endsWith(".class") && entryName.contains(basePackageName)) {
                                System.out.println("before" + Arrays.toString(bytes));
                                bytes = Dencrypt.encrypt(bytes,bindMachine);
                                System.out.println("after " + Arrays.toString(bytes));
                                break;
                            }
                        }

                        // jar中内嵌jar以外的普通文件一般是使用DEFLATE
                        // 要是jarEntry没有显式设定method也便是压缩方式,调用putNextEntry时候会设置成DEFLTE
                        JarEntry jarEntry = new JarEntry(entry_.getValue().getName());

                        // 要是使用STORED不压缩也是可以
                       /* jarEntry.setMethod(ZipEntry.STORED);
                        jarEntry.setSize(bytes.length);
                        jarEntry.setCompressedSize(bytes.length);
                        CRC32 crc32 = new CRC32();
                        crc32.update(bytes);
                        jarEntry.setCrc(crc32.getValue());*/

                        synchronized (encryptedJar) {
                            encryptedJar.putNextEntry(jarEntry);
                            encryptedJar.write(bytes);
                            encryptedJar.closeEntry();
                        }

                        byteArrayOutputStream.reset();
                    }
                }
            }
        }
    }


   /* private static void encrypt(byte[] orginalByteArr) {

        if (StringUtils.isBlank(System.getProperty(""))) {
            Dencrypt.encrypt(orginalByteArr, bindMachine);
            return;
        }


        orginalByteArr[0] = (byte) 0x01;
        if (bindMachine) {
            orginalByteArr[1] = (byte) 0xdd;
        } else {
            orginalByteArr[1] = (byte) 0x00;
        }
        orginalByteArr[2] = (byte) 0x01;
        orginalByteArr[3] = (byte) 0x00;
        // 0100 0000
        for (int a = 4; a < orginalByteArr.length; a++) {
            int low = orginalByteArr[a] & 0x0f;
            int high = (orginalByteArr[a] >> 4) & 0x0f;
            orginalByteArr[a] = (byte) ((low << 4) | high);
        }

    }*/

    private static class Pair {
        // 代表当前没有处在嵌套的jar
        public static final Pair EMPTY = new Pair("", null);

        public Pair(String subJarGlobalEntryName, JarEntry jarEntry) {
            this.subJarGlobalEntryName = subJarGlobalEntryName;
            this.jarEntry = jarEntry;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            Pair pair = (Pair) o;
            return Objects.equals(subJarGlobalEntryName, pair.subJarGlobalEntryName);
        }

        @Override
        public int hashCode() {
            return Objects.hash(subJarGlobalEntryName);
        }

        public String subJarGlobalEntryName;
        public JarEntry jarEntry;
    }

    private static void waitDone(List<Future<Void>> futureList) throws Exception {
        for (Future<Void> future : futureList) {
            future.get(90, TimeUnit.SECONDS);
        }
        futureList.clear();
    }

    private static void rm(String path) {
        File file = new File(path);

        List<Future> futureList = new ArrayList<>();
        if (file.isDirectory()) {
            File[] files = file.listFiles();

            if (null == files || files.length == 0) {
                file.delete();
                return;
            }

            for (File file0 : files) {
                Runnable runnable = () -> rm(file0.getPath());
                if (Thread.currentThread().getName().equals("main")) {
                    Future future = POOL.submit(runnable);
                    futureList.add(future);
                } else {
                    runnable.run();
                }

            }
        }

        for (Future future : futureList) {
            try {
                future.get();
            } catch (Exception e) {
                //
            }
        }

        file.delete();
    }
}
